/****************************************************************
** Product  :   HP Subscriber Network Application Policy
** Module   :   gr-manager
** Date: Apr 10, 2015               
** Author: Joey Yi
** (C) Copyright 2015, Hewlett-Packard Company, All Rights Reserved.
** This software is the proprietary information of HP, Inc.  
** Use is subject to license terms.
*****************************************************************
** Description:   
**  SSHTools encapsulate a set of methods running over the SSH. 
**     Execute remote command through by SSH;
**     Get a command result code;
**     Copy local files to remote host;
**     Copy remote files to local host;
****************************************************************/
package com.hp.snap.gr.utils;

import com.jcraft.jsch.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.util.FileCopyUtils;

import java.io.*;
import java.util.HashMap;
import java.util.Map;

/**
 * SSHTools encapsulate a set of methods running over the SSH. 
 *     Execute remote command through by SSH;
 *     Get a command result code;
 *     Copy local files to remote host;
 *     Copy remote files to local host; 
 *
 * <P>
 * Version: 3.3 <br>
 * Author: Joey Yi
 * Date: Apr 10, 2015
 * 
 * </P>
 *
 **/
public class SSHTools {

    private static Logger logger = LoggerFactory.getLogger(SSHTools.class);
    static private Map<String, Session> _mapCacheHostUser2Session = new HashMap<String, Session>();
    private final String m_username;
    private final String m_keyFile;
    //	private int _timeout = 0;//5000; //TODO sometimes cannot connect if timeout is set.
    private int _timeout = 50000;
    private ProxyHTTP _proxyHttp;
    private UserInfo _proxyUser;

    public SSHTools(String username, String key) {
        // If a username is not specified, default to the user who started the process.
        if (username == null || username.isEmpty()) {
            m_username = System.getProperty("user.name");
        } else {
            m_username = username;
        }
        // If a private key is not specified, default to the rsa key in the default location.
        if (key == null || key.isEmpty()) {
            m_keyFile = System.getProperty("user.home") + "/.ssh/id_rsa";
        } else {
            m_keyFile = key;
        }
        //TODO sometimes cannot connect if timeout is set.
//		if (timeout < 3000) throw new IllegalArgumentException("SSH timeout must be great than 1 second.");
//		_timeout = timeout;

    }

    static int checkAck(InputStream in) throws IOException {
        int b = in.read();
        // b may be 0 for success,
        //          1 for error,
        //          2 for fatal error,
        //          -1
        if (b == 0) return b;
        if (b == -1) return b;

        if (b == 1 || b == 2) {
            StringBuffer sb = new StringBuffer();
            int c;
            do {
                c = in.read();
                sb.append((char) c);
            }
            while (c != '\n');
            if (b == 1) { // error
                logger.error(sb.toString());
            }
            if (b == 2) { // fatal error
                logger.error(sb.toString());
            }
        }
        return b;
    }

    public static String stringify(String[] str_array) {
        StringBuffer result = new StringBuffer(2048);
        if (str_array != null) {
            for (int i = 0; i < str_array.length; i++) {
                result.append(" ").append(str_array[i]);
            }
        }
        return result.toString();
    }

    public int getTimeout() {
        return _timeout;
    }

    // Execute a remote SSH command on the specified host.
    public String cmd(String hostname, String command) throws Exception {
        return cmdSSH(m_username, m_keyFile, hostname, command);
    }

    // Execute a remote SSH command on the specified host.
    public String cmd(String hostname, String[] command) throws Exception {
        return cmdSSH(m_username, m_keyFile, hostname, command);
    }

    public String cmdSSH(String user, String key, String host, String[] command) throws Exception {
        return cmdSSH(user, key, host, stringify(command));
    }

    public String cmdSSH(String user, String key, String host, String command) throws Exception {
        Session session = getSession(user, key, host);

        String result = sshCmdOnSession(session, command);

        return result;
    }

    private String getSessionKey(String user, String host) {
        return Thread.currentThread().getId() + "-" + user + "@" + host;
        //return user + "@" + host;
    }

    private Session getSession(String user, String key, String host) throws JSchException {
        synchronized (_mapCacheHostUser2Session) {
            String sessionKey = getSessionKey(user, host);
            Session session = _mapCacheHostUser2Session.get(sessionKey);
            if ((session != null) && session.isConnected()) return session;
            if (logger.isDebugEnabled()) logger.debug("Connecting to SSH session by " + sessionKey);
            JSch jsch = new JSch();
            // Set the private key
            if (null != key)
                jsch.addIdentity(key);
            session = jsch.getSession(user, host, 22);
            session.setTimeout(_timeout);
            //Http Proxy?
            if (_proxyHttp != null) session.setProxy(_proxyHttp);
            if (_proxyUser != null) session.setUserInfo(_proxyUser);

            // To avoid the UnknownHostKey issue
            java.util.Properties config = new java.util.Properties();
            config.put("StrictHostKeyChecking", "no");
            session.setConfig(config);
            session.connect();
            Session old = _mapCacheHostUser2Session.put(sessionKey, session);
            if (old != null) old.disconnect();
            return session;
        }
    }

    public void closeAll() {
        synchronized (_mapCacheHostUser2Session) {
            for (Session s : _mapCacheHostUser2Session.values()) {
                s.disconnect();
            }
        }
    }

    public boolean copyFromLocal(String src, String hostNameTo, String pathTo) throws Exception {
        return ScpTo(src, m_username, m_keyFile, hostNameTo, pathTo);
    }

    public boolean copyFromRemote(String dst, String hostNameFrom, String pathFrom) {
        return ScpFrom(m_username, m_keyFile, hostNameFrom, pathFrom, dst);
    }

    // The Jsch method for SCP to.
    // This code is direcly copied from the Jsch SCP sample program.
    public boolean ScpTo(String local_file, String user, String key, String host, String remote_file) throws Exception {

        FileInputStream fis = null;
        try {
            boolean ptimestamp = true;
            String command = "scp " + (ptimestamp ? "-p" : "") + " -t " + remote_file;
            if (logger.isDebugEnabled()) logger.debug("CMD: '" + CmdUtils.printableCommand(command) + "'");

            JSch jsch = new JSch();
            // Set the private key
            if (null != key)
                jsch.addIdentity(key);
            Session session = jsch.getSession(user, host, 22);

            // To avoid the UnknownHostKey issue
            java.util.Properties config = new java.util.Properties();
            config.put("StrictHostKeyChecking", "no");
            session.setConfig(config);
            //Http Proxy?
            if (_proxyHttp != null) session.setProxy(_proxyHttp);
            if (_proxyUser != null) session.setUserInfo(_proxyUser);

            session.connect(_timeout); // timeout after 5 seconds

            // exec 'scp -t rfile' remotely
            Channel channel = session.openChannel("exec");
            ((ChannelExec) channel).setCommand(command);

            // get I/O streams for remote scp
            OutputStream out = channel.getOutputStream();
            InputStream in = channel.getInputStream();

            channel.connect();

            if (checkAck(in) != 0) {
                return false;
            }

            File _lfile = new File(local_file);

            if (ptimestamp) {
                command = "T " + (_lfile.lastModified() / 1000) + " 0";
                // The access time should be sent here,
                // but it is not accessible with JavaAPI ;-<
                command += (" " + (_lfile.lastModified() / 1000) + " 0\n");
                out.write(command.getBytes());
                out.flush();
                if (checkAck(in) != 0) {
                    return false;
                }
            }

            // send "C0644 filesize filename", where filename should not include '/'
            long filesize = _lfile.length();
            command = "C0644 " + filesize + " ";
            if (local_file.lastIndexOf('/') > 0) {
                command += local_file.substring(local_file.lastIndexOf('/') + 1);
            } else {
                command += local_file;
            }
            command += "\n";
            out.write(command.getBytes());
            out.flush();
            if (checkAck(in) != 0) {
                return false;
            }

            // send a content of lfile
            fis = new FileInputStream(local_file);
            byte[] buf = new byte[1024];
            while (true) {
                int len = fis.read(buf, 0, buf.length);
                if (len <= 0) break;
                out.write(buf, 0, len); //out.flush();
            }
            fis.close();
            fis = null;
            // send '\0'
            buf[0] = 0;
            out.write(buf, 0, 1);
            out.flush();
            if (checkAck(in) != 0) {
                return false;
            }
            out.close();

            channel.disconnect();
            session.disconnect();
        } finally {
            try {
                if (fis != null) fis.close();
            } catch (Exception ee) {
            }
        }

        return true;
    }

    // The Jsch method for SCP from.
    // This code is directly copied from the Jsch SCP sample program.  Error handling has been modified by VoltDB.
    public boolean ScpFrom(String user, String key, String host, String remote_file, String local_file) {

        FileOutputStream fos = null;
        try {
            String prefix = null;
            if (new File(local_file).isDirectory()) {
                prefix = local_file + File.separator;
            }

            String command = "scp -f " + remote_file;
            if (logger.isDebugEnabled()) logger.debug("CMD: '" + CmdUtils.printableCommand(command) + "'");

            JSch jsch = new JSch();
            // Set the private key
            if (null != key)
                jsch.addIdentity(key);
            Session session = jsch.getSession(user, host, 22);

            // To avoid the UnknownHostKey issue
            java.util.Properties config = new java.util.Properties();
            config.put("StrictHostKeyChecking", "no");
            session.setConfig(config);

            //Http Proxy?
            if (_proxyHttp != null) session.setProxy(_proxyHttp);
            if (_proxyUser != null) session.setUserInfo(_proxyUser);

            session.connect();

            // exec 'scp -f rfile' remotely
            Channel channel = session.openChannel("exec");
            ((ChannelExec) channel).setCommand(command);

            // get I/O streams for remote scp
            OutputStream out = channel.getOutputStream();
            InputStream in = channel.getInputStream();

            channel.connect();
            byte[] buf = new byte[1024];

            // send '\0'
            buf[0] = 0;
            out.write(buf, 0, 1);
            out.flush();

            while (true) {
                int c = checkAck(in);
                if (c != 'C') {
                    break;
                }

                // read '0644 '
                in.read(buf, 0, 5);

                long filesize = 0L;
                while (true) {
                    if (in.read(buf, 0, 1) < 0) {
                        // error
                        break;
                    }
                    if (buf[0] == ' ') break;
                    filesize = filesize * 10L + (buf[0] - '0');
                }

                String file = null;
                for (int i = 0; ; i++) {
                    in.read(buf, i, 1);
                    if (buf[i] == (byte) 0x0a) {
                        file = new String(buf, 0, i);
                        break;
                    }
                }

                String destination_file = prefix == null ? local_file : prefix + file;
                if (logger.isDebugEnabled()) logger.debug("CMD: scp to local file '" + destination_file + "'");

                // send '\0'
                buf[0] = 0;
                out.write(buf, 0, 1);
                out.flush();

                // read a content of lfile
                fos = new FileOutputStream(destination_file);
                int foo;
                while (true) {
                    if (buf.length < filesize) foo = buf.length;
                    else foo = (int) filesize;
                    foo = in.read(buf, 0, foo);
                    if (foo < 0) {
                        // error
                        break;
                    }
                    fos.write(buf, 0, foo);
                    filesize -= foo;
                    if (filesize == 0L) break;
                }
                fos.close();
                fos = null;

                if (checkAck(in) != 0) {
                    if (logger.isDebugEnabled()) logger.debug("CMD: scp checkAck failed");
                    System.out.println("checkAck did not equal zero.");
                    return false;
                }

                // send '\0'
                buf[0] = 0;
                out.write(buf, 0, 1);
                out.flush();
            }

            session.disconnect();
        } catch (Exception e) {
            logger.error(e.toString());
            if (logger.isDebugEnabled()) logger.debug("CMD: scp failed with exception: " + e.toString());
            return false;
        } finally {
            try {
                if (fos != null) fos.close();
            } catch (Exception ee) {
            }
        }

        return true;
    }

    private String sshCmdOnSession(Session session, String command) throws JSchException, IOException {
        StringBuilder result = new StringBuilder(2048);
        Channel channel = session.openChannel("exec");
        ((ChannelExec) channel).setCommand(command);
        // Direct stderr output of command
        InputStream err = ((ChannelExec) channel).getErrStream();
        InputStreamReader errStrRdr = new InputStreamReader(err, "UTF-8");
        Reader errStrBufRdr = new BufferedReader(errStrRdr);

        // Direct stdout output of command
        InputStream out = channel.getInputStream();
        InputStreamReader outStrRdr = new InputStreamReader(out, "UTF-8");
        Reader outStrBufRdr = new BufferedReader(outStrRdr);

        StringBuffer stdout = new StringBuffer();
        StringBuffer stderr = new StringBuffer();
        int exitStatus = -1;

        channel.connect(_timeout);  // timeout after 5 seconds
        while (true) {
            if (channel.isClosed()) {
                exitStatus = channel.getExitStatus();
                break;
            }

            // Read from both streams here so that they are not blocked,
            // if they are blocked because the buffer is full, channel.isClosed() will never
            // be true.
            int ch;
            while (outStrBufRdr.ready() && (ch = outStrBufRdr.read()) > -1) {
                stdout.append((char) ch);
            }
            while (errStrBufRdr.ready() && (ch = errStrBufRdr.read()) > -1) {
                stderr.append((char) ch);
            }

            try {
                Thread.sleep(100);
            } catch (InterruptedException ie) {
            }
        }

        // In case there's still some more stuff in the buffers, read them
        int ch;
        while ((ch = outStrBufRdr.read()) > -1) {
            stdout.append((char) ch);
        }
        while ((ch = errStrBufRdr.read()) > -1) {
            stderr.append((char) ch);
        }

        outStrBufRdr.close();
        errStrBufRdr.close();
        // After the command is executed, gather the results (both stdin and stderr).
        result.append(exitStatus).append('|');
        result.append(stdout.toString()).append("|");
        result.append(stderr.toString());

        // Shutdown the connection
        channel.disconnect();

        if (logger.isTraceEnabled()) {
            logger.trace("executed command {} on {}, result is {}", new Object[]{CmdUtils.printableCommand(command), session.getHost(), result});
        }
        return result.toString();
    }

    /**
     * Get exit status of command result output
     *
     * @param commandResult Integer.MIN-uinknown, others code
     * @return
     */
    public int getExitStatusOfCommandResult(String commandResult) {
        if (Utils.isEmpty(commandResult)) return 0;//success
        int result = Integer.MIN_VALUE;
        int i = commandResult.indexOf('|');
        if (i >= 0) {
            result = Integer.parseInt(commandResult.substring(0, i));
        }
        return result;
    }

    public String getCommandErr(String commandResult) {
        String[] vars = commandResult.split("\\s*\\|\\s*");
        if (vars.length > 2) {
            return vars[2];
        } else {
            if (logger.isTraceEnabled()) {
                logger.trace("command result in illegal format: {}", commandResult);
            }
            return "";
        }
    }

    public String getCommandOutput(String commandResult) {
        String[] vars = commandResult.split("\\s*\\|\\s*");
        if (vars.length > 1) {
            return vars[1];
        } else {
            if (logger.isTraceEnabled()) {
                logger.trace("command result in illegal format: {}", commandResult);
            }
            return "";
        }
//        String result = "";
//        int i = commandResult.indexOf('|');
//        if (i >= 0) {
//            result = commandResult.substring(i + 1);
//            result = result.trim();
//        }
//        return result;
    }

    public void setHttpProxy(String address, int port) {
        if (address == null) _proxyHttp = null;
        _proxyHttp = new ProxyHTTP(address, port);
    }

    public void setProxyUserInfo(final String user, final String password) {
        if (user == null) _proxyUser = null;
        _proxyUser = new UserInfo() {
            @Override
            public String getPassphrase() {
                return user;
            }

            @Override
            public String getPassword() {
                return password;
            }

            @Override
            public boolean promptPassword(String s) {
                return false;
            }

            @Override
            public boolean promptPassphrase(String s) {
                return false;
            }

            @Override
            public boolean promptYesNo(String s) {
                return false;
            }

            @Override
            public void showMessage(String s) {

            }
        };

    }

    protected File genCmdFile(String content) throws IOException {
        String filecontent = ". ~/.bash_profile\n" + content;
        File rs = File.createTempFile("REMOTE", ".sh", new File(Constants.USER_GR_TEMP));
        OutputStreamWriter osWriter = new FileWriter(rs);
        try {
            FileCopyUtils.copy(filecontent, osWriter);
        } finally {
            if (osWriter != null)
                osWriter.close();
        }
        return rs;
    }

    public String executeRemoteShell(boolean _isFile, String content, String toHost) throws Exception {
        //if(_log.isDebugEnabled()) _log.debug(_hostName + "(" + _hostTo + ")>Deploying command " + _commandName + ".");
        File _fileCmdLocal = genCmdFile(content);
        StringBuilder result = new StringBuilder();
        //Copy file
        String remotePath = Utils.getConfigurationValue("USER_GR_TEMP");
        String _remoteCmdFile = remotePath + "/" + _fileCmdLocal.getName();
        String mkdirResult = cmd(toHost, "mkdir -p " + remotePath);

//		if(logger.isDebugEnabled()){
//			logger.debug("Create remote path {}, result is {}", new Object[]{remotePath,mkdirResult});
//		}
        copyFromLocal(_fileCmdLocal.getAbsolutePath(), toHost, _remoteCmdFile);
        cmd(toHost, "chmod +x " + _remoteCmdFile);
        //return cmd(toHost, _remoteCmdFile);
        String resultCmd = cmd(toHost, _remoteCmdFile);
        cmd(toHost, "rm " + _remoteCmdFile);
        _fileCmdLocal.delete();
        return resultCmd;
        //
        //_remoteCmdFile = remotePath + "/" + _fileCmdLocal.getName();
        //sshCopyFromLocal(_fileCmdLocal.getAbsolutePath(), _remoteCmdFile, result);
//		sshCmd("chmod +x " + _remoteCmdFile, result, true);
//		cmd(toHost, "mkdir -p " + remotePath);
    }
    
    public String executeRemoteScriptFile(File scritFile, String scriptCommand, String toHost) throws Exception {
        String remotePath = "/tmp";
        String _remoteCmdFile = remotePath + "/" + scritFile.getName();
        copyFromLocal(scritFile.getAbsolutePath(), toHost, _remoteCmdFile);
        cmd(toHost, "chmod +x " + _remoteCmdFile);
        
        File _fileCmdLocal = genCmdFileInTMP(scriptCommand);
        String _remoteExecuteCmdFile = remotePath + "/" + _fileCmdLocal.getName();
        copyFromLocal(_fileCmdLocal.getAbsolutePath(), toHost, _remoteExecuteCmdFile);
        cmd(toHost, "chmod +x " + _remoteExecuteCmdFile);
        String resultCmd = cmd(toHost, _remoteExecuteCmdFile);
        cmd(toHost, "rm " + _remoteExecuteCmdFile);
        cmd(toHost, "rm " + _remoteCmdFile);
        return resultCmd;
        
    }
    
    protected File genCmdFileInTMP(String content) throws IOException {
        String filecontent = ". ~/.bash_profile\n" + content;
        File rs = File.createTempFile("REMOTE", ".sh", new File("/tmp"));
        OutputStreamWriter osWriter = new FileWriter(rs);
        try {
            FileCopyUtils.copy(filecontent, osWriter);
        } finally {
            if (osWriter != null)
                osWriter.close();
        }
        return rs;
    }

}
